import logging
import time
from ruamel import yaml


class AverageMeter(object):
    """用于计算和存储当前值、平均值、总和和计数"""

    def __init__(self):
        # 初始化时重置所有值
        self.reset()

    def reset(self):
        # 重置当前值、平均值、总和和计数
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        """更新当前值、总和、计数和平均值

        :param val: 当前批次的值
        :param n: 批次大小，默认为1
        """
        self.val = val
        self.sum += val * n
        self.count += n
        # 计算平均值
        self.avg = self.sum / self.count


def Prepare_logger(args, eval=False):
    """准备日志记录器

    :param args: 包含日志路径等信息的参数对象
    :param eval: 是否为评估模式，默认为False
    :return: 配置好的日志记录器
    """
    # 创建日志记录器
    logger = logging.getLogger(__name__)
    # 防止日志信息重复输出
    logger.propagate = False
    # 设置日志级别为INFO
    logger.setLevel(logging.INFO)

    # 创建控制台处理器，并设置日志格式
    handler = logging.StreamHandler()
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    handler.setFormatter(formatter)
    # 设置处理器级别
    handler.setLevel(0)
    # 将处理器添加到日志记录器
    logger.addHandler(handler)

    # 获取当前时间并格式化为字符串
    date = time.strftime('%Y%m%d%H%M', time.localtime(time.time()))
    # 根据是否为评估模式决定日志文件名
    logfile = args.snapshot_pref + date + '.log' if not eval else args.snapshot_pref + f'/{date}-Eval.log'

    # 创建文件处理器，并设置日志格式
    file_handler = logging.FileHandler(logfile, mode='w')
    file_handler.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    file_handler.setFormatter(formatter)
    # 将文件处理器添加到日志记录器
    logger.addHandler(file_handler)

    return logger


def get_configs(dataset):
    """从yaml文件中读取指定数据集的配置

    :param dataset: 数据集名称
    :return: 指定数据集的配置字典
    """
    # 读取配置文件
    with open('./configs/dataset_cfg.yaml', 'r') as file:
        data = yaml.load(file)
    # 返回指定数据集的配置
    return data[dataset]


def get_and_save_args(parser):
    """解析命令行参数，更新默认配置，并保存更新后的配置到yaml文件

    :param parser: ArgumentParser对象，用于解析命令行参数
    :return: 更新后的默认配置字典
    """
    # 解析命令行参数
    args = parser.parse_args()

    # 加载默认配置文件
    with open('./configs/default_config.yaml', 'r') as file:
        default_config = yaml.load(file, Loader=yaml.RoundTripLoader)

    # 获取命令行参数的字典形式
    current_config = vars(args)

    # 遍历当前参数，如果与默认配置不同则更新默认配置
    for k, v in current_config.items():
        if k in default_config:
            if (v != default_config[k]) and (v is not None):
                print(f"Updating:  {k}: {default_config[k]} (default) ----> {v}")
                default_config[k] = v

    # 将更新后的默认配置保存到新的yaml文件
    with open('./current_configs.yaml', 'w') as file:
        yaml.dump(default_config, file, indent=4, Dumper=yaml.RoundTripDumper)

    return default_config
